"""
收敛性检查模块

用于检查数值计算是否达到收敛
"""

import numpy as np
from .. import config as cfg

class ConvergenceChecker:
    """收敛性检查器"""
    
    def __init__(self, tolerance=cfg.CONVERGENCE_TOLERANCE):
        """
        初始化收敛性检查器
        
        参数：
        tolerance: float, 收敛判据
        """
        self.tolerance = tolerance
        
    def check_field_convergence(self, old_field, new_field):
        """
        检查场变量是否收敛
        
        参数：
        old_field: array, 上一时间步的场变量
        new_field: array, 当前时间步的场变量
        
        返回：
        bool, 是否收敛
        float, 最大相对误差
        """
        # 计算相对误差
        diff = np.abs(new_field - old_field)
        max_val = np.maximum(np.abs(new_field), np.abs(old_field))
        max_val = np.where(max_val > 0, max_val, 1.0)  # 避免除以零
        relative_error = diff / max_val
        
        # 获取最大相对误差
        max_error = np.max(relative_error)
        
        return max_error < self.tolerance, max_error
        
    def check_mass_conservation(self, channels):
        """
        检查质量守恒
        
        参数：
        channels: dict, 子通道字典
        
        返回：
        bool, 是否满足质量守恒
        float, 最大质量不平衡
        """
        max_imbalance = 0.0
        
        for channel in channels.values():
            # 计算入口质量流量
            props_in = channel.fluid_props.get_properties(
                channel.pressure[0], channel.T_l[0], 'liquid')
            if props_in is None:
                continue
                
            m_in = props_in['density'] * channel.u_l[0] * channel.area
            
            # 计算出口质量流量
            props_out_l = channel.fluid_props.get_properties(
                channel.pressure[-1], channel.T_l[-1], 'liquid')
            props_out_v = channel.fluid_props.get_properties(
                channel.pressure[-1], channel.T_v[-1], 'vapor')
                
            if props_out_l is None or props_out_v is None:
                continue
                
            m_out = (channel.alpha_l[-1] * props_out_l['density'] * channel.u_l[-1] +
                    channel.alpha_v[-1] * props_out_v['density'] * channel.u_v[-1]) * channel.area
            
            # 计算相对不平衡
            imbalance = abs(m_in - m_out) / m_in
            max_imbalance = max(max_imbalance, imbalance)
            
        return max_imbalance < self.tolerance, max_imbalance
        
    def check_energy_conservation(self, channels):
        """
        检查能量守恒
        
        参数：
        channels: dict, 子通道字典
        
        返回：
        bool, 是否满足能量守恒
        float, 最大能量不平衡
        """
        max_imbalance = 0.0
        
        for channel in channels.values():
            # 计算入口能量流率
            props_in = channel.fluid_props.get_properties(
                channel.pressure[0], channel.T_l[0], 'liquid')
            if props_in is None:
                continue
                
            e_in = props_in['density'] * channel.u_l[0] * props_in['enthalpy'] * channel.area
            
            # 计算出口能量流率
            props_out_l = channel.fluid_props.get_properties(
                channel.pressure[-1], channel.T_l[-1], 'liquid')
            props_out_v = channel.fluid_props.get_properties(
                channel.pressure[-1], channel.T_v[-1], 'vapor')
                
            if props_out_l is None or props_out_v is None:
                continue
                
            e_out = (channel.alpha_l[-1] * props_out_l['density'] * channel.u_l[-1] * 
                    props_out_l['enthalpy'] +
                    channel.alpha_v[-1] * props_out_v['density'] * channel.u_v[-1] * 
                    props_out_v['enthalpy']) * channel.area
            
            # 计算加热功率
            q_total = cfg.MAX_HEAT_FLUX * cfg.CHANNEL_LENGTH
            
            # 计算相对不平衡
            imbalance = abs(e_out - e_in - q_total) / e_in
            max_imbalance = max(max_imbalance, imbalance)
            
        return max_imbalance < self.tolerance, max_imbalance
        
    def check_convergence(self, channels, old_channels):
        """
        检查整体收敛性
        
        参数：
        channels: dict, 当前时间步的子通道字典
        old_channels: dict, 上一时间步的子通道字典
        
        返回：
        bool, 是否收敛
        dict, 包含各项检查结果的字典
        """
        results = {}
        
        # 检查场变量收敛性
        for channel_id, channel in channels.items():
            old_channel = old_channels[channel_id]
            
            # 检查液相分数
            conv_alpha_l, err_alpha_l = self.check_field_convergence(
                old_channel.alpha_l, channel.alpha_l)
            results[f'alpha_l_channel_{channel_id}'] = {
                'converged': conv_alpha_l,
                'error': err_alpha_l
            }
            
            # 检查液相速度
            conv_u_l, err_u_l = self.check_field_convergence(
                old_channel.u_l, channel.u_l)
            results[f'u_l_channel_{channel_id}'] = {
                'converged': conv_u_l,
                'error': err_u_l
            }
            
            # 检查气相速度
            conv_u_v, err_u_v = self.check_field_convergence(
                old_channel.u_v, channel.u_v)
            results[f'u_v_channel_{channel_id}'] = {
                'converged': conv_u_v,
                'error': err_u_v
            }
            
            # 检查压力
            conv_p, err_p = self.check_field_convergence(
                old_channel.pressure, channel.pressure)
            results[f'pressure_channel_{channel_id}'] = {
                'converged': conv_p,
                'error': err_p
            }
            
            # 检查温度
            conv_t_l, err_t_l = self.check_field_convergence(
                old_channel.T_l, channel.T_l)
            results[f'T_l_channel_{channel_id}'] = {
                'converged': conv_t_l,
                'error': err_t_l
            }
            
        # 检查质量守恒
        mass_conv, mass_imb = self.check_mass_conservation(channels)
        results['mass_conservation'] = {
            'converged': mass_conv,
            'imbalance': mass_imb
        }
        
        # 检查能量守恒
        energy_conv, energy_imb = self.check_energy_conservation(channels)
        results['energy_conservation'] = {
            'converged': energy_conv,
            'imbalance': energy_imb
        }
        
        # 判断整体收敛性
        all_converged = all(item['converged'] for item in results.values())
        
        return all_converged, results
        
    def print_convergence_status(self, results):
        """
        打印收敛性检查结果
        
        参数：
        results: dict, 收敛性检查结果
        """
        print("\n收敛性检查结果：")
        print("-" * 50)
        
        # 打印场变量收敛性
        print("\n场变量收敛性：")
        for key, value in results.items():
            if 'channel' in key:
                status = "收敛" if value['converged'] else "未收敛"
                print(f"{key}: {status} (误差: {value['error']:.2e})")
                
        # 打印守恒性
        print("\n守恒性检查：")
        mass_status = "满足" if results['mass_conservation']['converged'] else "不满足"
        print(f"质量守恒: {mass_status} (不平衡: {results['mass_conservation']['imbalance']:.2e})")
        
        energy_status = "满足" if results['energy_conservation']['converged'] else "不满足"
        print(f"能量守恒: {energy_status} (不平衡: {results['energy_conservation']['imbalance']:.2e})")
        
        # 打印整体状态
        all_converged = all(item['converged'] for item in results.values())
        print("\n整体状态：", "收敛" if all_converged else "未收敛")
        print("-" * 50) 